import sys,os

try:
    import matplotlib.cbook as cbook
    from matplotlib.backend_bases import NavigationToolbar2, cursors
    from matplotlib import rcParams


    from matplotlib.backends.backend_wx import NavigationToolbar2Wx
    
    from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg
    
    from matplotlib.transforms import Bbox

except:
    print >>sys.stderr, "iplot.wx_toolbar requires backend_wx be installed"
    sys.exit()

import wx
from iplot.mpl.debug import DEBUG_MSG
from iplot.mpl.wx_utils import bind, load_bitmap
from iplot.mpl.tools import mask as _mask, unmask as _unmask

mplimages = os.path.join(rcParams['datapath'],'images')
thisdir  = os.path.dirname(__file__)
iplotimages = os.path.join(thisdir,'images')
_load_bitmap = lambda filename: load_bitmap(filename,[mplimages,iplotimages])

#from numpy import ma

class IPlotToolbar(NavigationToolbar2Wx):
    
    def __init__(self, parent, canvas):
        DEBUG_MSG("__init__", 1, self)
        wx.ToolBar.__init__(self, parent, -1)
        NavigationToolbar2.__init__(self, canvas)
        canvas.toolbar = self
        self._idle = True
        self.statbar = None
        
        self.iplot = self._parent.iplot
        self.datadlg = self._parent.iplot.datadlg
        self.propertydlg = self._parent.iplot.propertydlg

    def get_canvas(self, frame, fig):
        #this creates a new canvas... 
        return FigureCanvasWxAgg(frame, -1, fig)
    
    def _init_toolbar(self):
        #NavigationToolbar2Wx._init_toolbar(self)

        DEBUG_MSG("_init_toolbar", 1, self)
        
        self._parent = self.canvas.GetParent()

        _NTB2_HOME    =wx.NewId()
        self._NTB2_BACK    =wx.NewId()
        self._NTB2_FORWARD =wx.NewId()
        self._TB_PAN     =wx.NewId()
        self._TB_ZOOM    =wx.NewId()
        _NTB2_SAVE    = wx.NewId()
        _NTB2_SUBPLOT    =wx.NewId()
        _TB_EDIT = wx.NewId()

        self.SetToolBitmapSize(wx.Size(24,24))

        self.AddSimpleTool(_NTB2_HOME, _load_bitmap('home.png'),
                           'Home', 'Reset original view')
        self.AddSimpleTool(self._NTB2_BACK, _load_bitmap('back.png'),
                           'Back', 'Back navigation view')
        self.AddSimpleTool(self._NTB2_FORWARD, _load_bitmap('forward.png'),
                           'Forward', 'Forward navigation view')
        # todo: get new bitmap
        self.AddCheckTool(self._TB_PAN, _load_bitmap('move.png'),
                           shortHelp='Pan',
                           longHelp='Pan with left, zoom with right')
        self.AddCheckTool(self._TB_ZOOM, _load_bitmap('zoom_to_rect.png'),
                           shortHelp='Zoom', longHelp='Zoom to rectangle')

        self.AddSeparator()
        self.AddSimpleTool(_NTB2_SUBPLOT, _load_bitmap('subplots.png'),
                           'Configure subplots', 'Configure subplot parameters')
        
        self.AddSimpleTool(_TB_EDIT, _load_bitmap('edit.xpm'),
                           'Edit', 'Edit Axes')

        self.AddSimpleTool(_NTB2_SAVE, _load_bitmap('filesave.png'),
                           'Save', 'Save plot contents to file')

        bind(self, wx.EVT_TOOL, self.home, id=_NTB2_HOME)
        bind(self, wx.EVT_TOOL, self.forward, id=self._NTB2_FORWARD)
        bind(self, wx.EVT_TOOL, self.back, id=self._NTB2_BACK)
        bind(self, wx.EVT_TOOL, self.zoom, id=self._TB_ZOOM)
        bind(self, wx.EVT_TOOL, self.pan, id=self._TB_PAN)
        bind(self, wx.EVT_TOOL, self.configure_subplot, id=_NTB2_SUBPLOT)
        bind(self, wx.EVT_TOOL, self.save, id=_NTB2_SAVE)
        bind(self, wx.EVT_TOOL, self.edit, id=_TB_EDIT)
        
        self.AddSeparator()
        self._TB_DATASHEET = wx.NewId()
        self._TB_MASK = wx.NewId()
        self._TB_UNMASK = wx.NewId()
        
        
        #basedir = '/Users/chaffra/Projects/LABPLOT/pyAnalysis/iplot/mpl'
        #basedir = os.path.join(basedir,'images')
        
        self._tb_mask = self.AddCheckTool(self._TB_MASK, _load_bitmap('mask.xpm'),
                          shortHelp='Mask', longHelp='Mask data points')
        bind(self,wx.EVT_TOOL,self.mask,id=self._TB_MASK)
        
        self._tb_unmask = self.AddCheckTool(self._TB_UNMASK, _load_bitmap('unmask.xpm'),
                          shortHelp='Unmask', longHelp='Unmask selected axes')
        bind(self,wx.EVT_TOOL,self.unmask,id=self._TB_UNMASK)
        
        self.AddSeparator()

        self._tb_data = self.AddCheckTool(self._TB_DATASHEET, _load_bitmap('data.xpm'),
                          shortHelp='Datasheet', longHelp='Datasheet of selected axes')
        bind(self,wx.EVT_TOOL,self.datasheet,id=self._TB_DATASHEET)            
        self.Realize()
        
    def edit(self,evt=None):
        fig = self.canvas.figure
        self.propertydlg.browser.SetRoot(fig,"Figure %d" % fig.number)
        val = self.propertydlg.ShowModal()
        
        if val == wx.ID_OK:
            pass
        else:
            pass

    def set_state(self,canvas,state={}):
        self.canvas = canvas
        self.canvas.toolbar = self
        
        self._views = state['_views'] if state.has_key('_views') else  cbook.Stack() # a dict from axes index to a list of view limits
        self._positions = state['_positions'] if state.has_key('_positions') else  cbook.Stack() # stack of subplot positions
        self._xypress = None  # the  location and axis info at the time of the press
        self._active = state['_active'] if state.has_key('_active') else None
        #self._lastCursor = None
        #self._idDrag=self.canvas.mpl_disconnect(self._idDrag)
        #self._idDrag=
        self.canvas.mpl_connect('motion_notify_event', self.mouse_move)
        self._button_pressed = None # determined by the button pressed at start
        self.set_history_buttons()
        self._idle = True
        self.mode = state['mode'] if state.has_key('mode') else ''
        self.set_message(self.mode)
        
        
    def get_state(self):
        state = dict()
        state['_views'] = self._views
        state['_positions'] = self._positions
        state['_active'] = self._active
        state['mode'] = self.mode
        return state
    
    def get_default_state(self):
        state = dict()
        state['_views'] = cbook.Stack()
        state['_positions'] = cbook.Stack()
        return state
        
    
    def datasheet(self, evt=None):
        fig = self.canvas.figure
        
        self.datadlg._refresh()
        self.datadlg.select_figure(fig)
        # this does not return until the dialog is closed.
        val = self.datadlg.ShowModal()
        
        if val == wx.ID_OK:
            pass
        else:
            pass
        
        self.ToggleTool(self._TB_DATASHEET,False)
        
    def draw(self):
        NavigationToolbar2Wx.draw(self)
        #self.iplot.GetSizer().Layout()
        
    def zoom(self, *args):
        self.ToggleTool(self._TB_MASK, False)
        self.ToggleTool(self._TB_PAN, False)
        self.ToggleTool(self._TB_UNMASK, False)       
        NavigationToolbar2.zoom(self,*args)
        
    def pan(self, *args):
        self.ToggleTool(self._TB_MASK, False)
        self.ToggleTool(self._TB_ZOOM, False)
        self.ToggleTool(self._TB_UNMASK, False)
        NavigationToolbar2.pan(self,*args)
        
    def unmask(self,evt=None):
        'activate unmask mode'
        self.ToggleTool(self._TB_PAN, False)
        self.ToggleTool(self._TB_ZOOM, False)
        self.ToggleTool(self._TB_MASK, False)
        
        if self._active == 'UNMASK':
            self._active = None
        else:
            self._active = 'UNMASK'
            
        if self._idPress is not None:
            self._idPress=self.canvas.mpl_disconnect(self._idPress)
            self.mode = ''

        if self._idRelease is not None:
            self._idRelease=self.canvas.mpl_disconnect(self._idRelease)
            self.mode = ''

        if  self._active:
            self._idPress = self.canvas.mpl_connect('button_press_event', self.press_unmask)
            self._idRelease = self.canvas.mpl_connect('button_release_event', self.release_unmask)
            self.mode = 'unmask rect'
            self.canvas.widgetlock(self)
        else:
            self.canvas.widgetlock.release(self)

        for a in self.canvas.figure.get_axes():
            a.set_navigate_mode(self._active)

        self.set_message(self.mode)
        
    def press_unmask(self,event):
        #print 'press_unmask'
        self.press_mask(event)
    
    def release_unmask(self,event):
        bbs = self._get_bboxes(event)
        #print bb

        if bbs:
            for bb,a in bbs:
                _unmask(a,bb)
                        
            self.draw()
            self._xypress = None
            self._button_pressed = None

            #self.push_current()
            self.release(event)

        
    def mask(self, evt=None):
        'activate mask mode'
        self.ToggleTool(self._TB_PAN, False)
        self.ToggleTool(self._TB_ZOOM, False)
        self.ToggleTool(self._TB_UNMASK, False)

        if self._active == 'MASK':
            self._active = None
        else:
            self._active = 'MASK'

        if self._idPress is not None:
            self._idPress=self.canvas.mpl_disconnect(self._idPress)
            self.mode = ''

        if self._idRelease is not None:
            self._idRelease=self.canvas.mpl_disconnect(self._idRelease)
            self.mode = ''

        if  self._active:
            
            self._idPress = self.canvas.mpl_connect('button_press_event', self.press_mask)
            self._idRelease = self.canvas.mpl_connect('button_release_event', self.release_mask)
            self.mode = 'mask rect'
            self.canvas.widgetlock(self)
        else:
            self.canvas.widgetlock.release(self)

        for a in self.canvas.figure.get_axes():
            a.set_navigate_mode(self._active)

        self.set_message(self.mode)
        
    
    def press_mask(self,event):
        'the press mouse button in mask to rect mode callback'
        
        if event.button == 1:
            self._button_pressed=1
        elif  event.button == 3:
            self._button_pressed=3
        else:
            self._button_pressed=None
            return

        x, y = event.x, event.y

        # push the current view to define home if stack is empty
        #if self._views.empty(): self.push_current()

        self._xypress=[]
        for i, a in enumerate(self.canvas.figure.get_axes()):
            if x is not None and y is not None and a.in_axes(event):
                self._xypress.append(( x, y, a, i, a.viewLim.frozen(), a.transData.frozen()))

        self.press(event)
    
    
    def _get_bboxes(self,event):
        if not self._xypress: return

        bbs = []
        for cur_xypress in self._xypress:
            x, y = event.x, event.y
            lastx, lasty, a, ind, lim, trans = cur_xypress

            # ignore singular clicks - 5 pixels is a threshold
            if abs(x-lastx)<5 or abs(y-lasty)<5:
                self._xypress = None
                self.release(event)
                self.draw()
                #print "singular click"
                return bbs

        #x0, y0, x1, y1 = lim.extents
            
            inverse = a.transData.inverted()
            lastx, lasty = inverse.transform_point( (lastx, lasty) )
            x, y = inverse.transform_point( (x, y) )
            
            bb = Bbox([[lastx,lasty],[x,y]])
            bbs.append((bb,a))
        
        return bbs
        
    def release_mask(self,event):
        'the release mouse button callback in mask rect mode'
        bbs = self._get_bboxes(event)
        #print bb

        if bbs:
            for bb,a in bbs:
                _mask(a,bb)
                        
            self.draw()
            self._xypress = None
            self._button_pressed = None

            #self.push_current()
            self.release(event)
    
    def mouse_move(self,event):
        #print self._active
        #print self._idDrag
        
        if not event.inaxes or not self._active:
            if self._lastCursor != cursors.POINTER:
                self.set_cursor(cursors.POINTER)
                self._lastCursor = cursors.POINTER
        else:
            if self._active in ['ZOOM','MASK','UNMASK']:
                if self._lastCursor != cursors.SELECT_REGION:
                    self.set_cursor(cursors.SELECT_REGION)
                    self._lastCursor = cursors.SELECT_REGION
                if self._xypress:
                    x, y = event.x, event.y
                    lastx, lasty, a, ind, lim, trans = self._xypress[0]
                    self.draw_rubberband(event, x, y, lastx, lasty)
            elif (self._active=='PAN' and
                  self._lastCursor != cursors.MOVE):
                self.set_cursor(cursors.MOVE)
                self._lastCursor = cursors.MOVE

        if event.inaxes and event.inaxes.get_navigate():

            try: s = event.inaxes.format_coord(event.xdata, event.ydata)
            except ValueError: pass
            except OverflowError: pass
            else:
                if len(self.mode):
                    self.set_message('%s, %s' % (self.mode, s))
                else:
                    self.set_message(s)
        else: self.set_message(self.mode)
        
        
        

#Work in progress...
class MainToolbar(IPlotToolbar):
    pass
